{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "!pip install struct-eqtable==0.1.0\n",
    "!pip install pypandoc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from struct_eqtable import build_model\n",
    "import pypandoc\n",
    "import time\n",
    "pypandoc.download_pandoc()\n",
    "from pypandoc import convert_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build model\n",
    "USE_CUDA = True\n",
    "MAX_TIME = 60 # maximum waiting time of model inference\n",
    "model = build_model(\"/'path to your project'/models/TabRec/StructEqTable\", max_new_tokens=4096, max_time=MAX_TIME)\n",
    "if USE_CUDA:\n",
    "    model = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model inference\n",
    "raw_image = Image.open(\"./demo_1.png\")\n",
    "latex_code_path = \"./demo_1_latex.txt\"\n",
    "markdown_path =  \"./demo_1_markdown.md\"\n",
    "\n",
    "start = time.time()           \n",
    "with torch.no_grad():\n",
    "    output = model(raw_image)\n",
    "end = time.time()\n",
    "\n",
    "latex_code = output[0]\n",
    "\n",
    "with open(latex_code_path, 'w') as f_latex:\n",
    "    f_latex.write(latex_code)\n",
    "\n",
    "if (end-start) > MAX_TIME:\n",
    "    print(f\"The model inference time exceeds the maximum waiting time {MAX_TIME} seconds, the result may be incomplete.\\nPlease increase the maximum waiting time.\")\n",
    "else:\n",
    "    markdown_code = convert_text(latex_code, 'md', format='latex')\n",
    "    html_code = convert_text(latex_code, 'html', format='latex')\n",
    "    print(f\"Table Markdown code:\\n{markdown_code}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  },
  "vscode": {
   "interpreter": {
    "hash": "2966c681669ffe51cf424a54485f53fc80fee578831e524e5659bc9f96b671ce"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
